#include <queue>
#include <vector>
#include <arm_neon.h>
#include <cstdlib>
#include <cmath>
#include <algorithm>
#include <numeric>
#include <random>

std::priority_queue<std::pair<float, uint32_t>> simd_pq(
    float* base, float* query, size_t base_number, size_t vecdim, size_t k, size_t rerank_k = 800) 
{
    std::priority_queue<std::pair<float, uint32_t>> q;

    constexpr size_t m = 4;
    constexpr size_t ks = 256;
    const size_t sub_dim = vecdim / m;

    static bool initialized = false;
    static std::vector<std::vector<float>> codebooks(m);
    static std::vector<std::vector<uint8_t>> pq_codes(m);

    if (!initialized) {
        // 初始化码本和编码
        codebooks.resize(m);
        pq_codes.resize(m);
        std::vector<std::mt19937> gen(m);
        
        for (size_t k = 0; k < m; ++k) {
            codebooks[k].resize(ks * sub_dim);
            pq_codes[k].resize(base_number);

            // 提取当前子空间的所有子向量
            std::vector<float> sub_vectors(base_number * sub_dim);
            for (size_t i = 0; i < base_number; ++i) {
                const float* src = base + i * vecdim + k * sub_dim;
                std::copy(src, src + sub_dim, &sub_vectors[i * sub_dim]);
            }

            // K-means聚类初始化
            std::vector<float> centroids(ks * sub_dim);
            std::vector<size_t> indices(base_number);
            std::iota(indices.begin(), indices.end(), 0);
            std::shuffle(indices.begin(), indices.end(), gen[k]);
            
            // 随机初始化中心
            for (size_t j = 0; j < ks; ++j)
                std::copy_n(&sub_vectors[indices[j] * sub_dim], sub_dim, &centroids[j * sub_dim]);

            // 简化的K-means迭代
            constexpr int MAX_ITER = 10;
            for (int iter = 0; iter < MAX_ITER; ++iter) {
                std::vector<float> new_centroids(ks * sub_dim, 0.0f);
                std::vector<size_t> counts(ks, 0);

                // 分配样本到最近中心
                for (size_t i = 0; i < base_number; ++i) {
                    float min_dist = std::numeric_limits<float>::max();
                    size_t best_j = 0;

                    for (size_t j = 0; j < ks; ++j) {
                        float dist = 0.0f;
                        const float* c = &centroids[j * sub_dim];
                        const float* v = &sub_vectors[i * sub_dim];

                        // NEON加速距离计算
                        float32x4_t sum = vdupq_n_f32(0.0f);
                        size_t d = 0;
                        for (; d + 4 <= sub_dim; d += 4) {
                            float32x4_t cv = vld1q_f32(v + d);
                            float32x4_t cc = vld1q_f32(c + d);
                            float32x4_t diff = vsubq_f32(cv, cc);
                            sum = vmlaq_f32(sum, diff, diff);
                        }
                        
                        // 处理剩余维度
                        float32x2_t sum2 = vadd_f32(vget_low_f32(sum), vget_high_f32(sum));
                        dist += vget_lane_f32(sum2, 0) + vget_lane_f32(sum2, 1);
                        for (; d < sub_dim; ++d)
                            dist += (v[d] - c[d]) * (v[d] - c[d]);

                        if (dist < min_dist) {
                            min_dist = dist;
                            best_j = j;
                        }
                    }

                    // 累加中心
                    const float* v = &sub_vectors[i * sub_dim];
                    float* nc = &new_centroids[best_j * sub_dim];
                    for (size_t d = 0; d < sub_dim; ++d)
                        nc[d] += v[d];
                    counts[best_j]++;
                }

                // 更新中心并处理空簇
                for (size_t j = 0; j < ks; ++j) {
                    if (counts[j] == 0) {
                        size_t idx = std::uniform_int_distribution<size_t>(0, base_number-1)(gen[k]);
                        std::copy_n(&sub_vectors[idx * sub_dim], sub_dim, &centroids[j * sub_dim]);
                    } else {
                        float scale = 1.0f / counts[j];
                        float* c = &centroids[j * sub_dim];
                        float* nc = &new_centroids[j * sub_dim];
                        for (size_t d = 0; d < sub_dim; d += 4) {
                            float32x4_t vec = vld1q_f32(nc + d);
                            vec = vmulq_n_f32(vec, scale);
                            vst1q_f32(c + d, vec);
                        }
                    }
                }
            }

            // 编码数据
            for (size_t i = 0; i < base_number; ++i) {
                float min_dist = std::numeric_limits<float>::max();
                size_t best_j = 0;
                const float* v = &sub_vectors[i * sub_dim];

                for (size_t j = 0; j < ks; ++j) {
                    const float* c = &centroids[j * sub_dim];
                    float dist = 0.0f;

                    float32x4_t sum = vdupq_n_f32(0.0f);
                    size_t d = 0;
                    for (; d + 4 <= sub_dim; d += 4) {
                        float32x4_t cv = vld1q_f32(v + d);
                        float32x4_t cc = vld1q_f32(c + d);
                        float32x4_t diff = vsubq_f32(cv, cc);
                        sum = vmlaq_f32(sum, diff, diff);
                    }

                    float32x2_t sum2 = vadd_f32(vget_low_f32(sum), vget_high_f32(sum));
                    dist += vget_lane_f32(sum2, 0) + vget_lane_f32(sum2, 1);
                    for (; d < sub_dim; ++d)
                        dist += (v[d] - c[d]) * (v[d] - c[d]);

                    if (dist < min_dist) {
                        min_dist = dist;
                        best_j = j;
                    }
                }
                pq_codes[k][i] = static_cast<uint8_t>(best_j);
            }

            // 存储码本
            std::copy(centroids.begin(), centroids.end(), codebooks[k].begin());
        }

        initialized = true;
    }

    // 预处理查询：计算各子空间的距离表
    std::vector<std::vector<float>> distance_tables(m, std::vector<float>(ks));
    for (size_t k = 0; k < m; ++k) {
        float* sub_query = query + k * sub_dim;
        float* codebook = codebooks[k].data();

        for (size_t j = 0; j < ks; ++j) {
            float* code = codebook + j * sub_dim;
            float sum = 0.0f;

            float32x4_t sum_vec = vdupq_n_f32(0.0f);
            size_t d;
            for (d = 0; d + 4 <= sub_dim; d += 4) {
                float32x4_t q = vld1q_f32(sub_query + d);
                float32x4_t c = vld1q_f32(code + d);
                float32x4_t diff = vsubq_f32(q, c);
                sum_vec = vmlaq_f32(sum_vec, diff, diff);
            }

            float32x2_t sum_low = vget_low_f32(sum_vec);
            float32x2_t sum_high = vget_high_f32(sum_vec);
            sum_low = vadd_f32(sum_low, sum_high);
            sum += vget_lane_f32(sum_low, 0) + vget_lane_f32(sum_low, 1);

            for (; d < sub_dim; ++d) {
                float diff = sub_query[d] - code[d];
                sum += diff * diff;
            }
            distance_tables[k][j] = sum;
        }
    }

    // 批量计算近似距离
    std::vector<float> approx_dists(base_number, 0.0f);
    constexpr size_t batch_size = 4;
    for (size_t i = 0; i + batch_size <= base_number; i += batch_size) {
        float32x4_t acc = vdupq_n_f32(0.0f);
        for (size_t k = 0; k < m; ++k) {
            uint8_t codes[batch_size];
            for (size_t b = 0; b < batch_size; ++b)
                codes[b] = pq_codes[k][i + b];

            float32x4_t dists;
            dists = vld1q_lane_f32(&distance_tables[k][codes[0]], dists, 0);
            dists = vld1q_lane_f32(&distance_tables[k][codes[1]], dists, 1);
            dists = vld1q_lane_f32(&distance_tables[k][codes[2]], dists, 2);
            dists = vld1q_lane_f32(&distance_tables[k][codes[3]], dists, 3);

            acc = vaddq_f32(acc, dists);
        }
        vst1q_f32(&approx_dists[i], acc);
    }

    // 处理剩余样本
    for (size_t i = base_number - (base_number % batch_size); i < base_number; ++i) {
        for (size_t k = 0; k < m; ++k)
            approx_dists[i] += distance_tables[k][pq_codes[k][i]];
    }

    // 使用小顶堆保留前rerank_k个候选
    using Queue = std::priority_queue<
        std::pair<float, uint32_t>,
        std::vector<std::pair<float, uint32_t>>,
        std::less<std::pair<float, uint32_t>>
    >;
    Queue temp_queue;

    for (size_t i = 0; i < base_number; ++i) {
        if (temp_queue.size() < rerank_k) {
            temp_queue.emplace(approx_dists[i], i);
        } else if (approx_dists[i] < temp_queue.top().first) {
            temp_queue.pop();
            temp_queue.emplace(approx_dists[i], i);
        }
    }

    // 精确重排序
    std::vector<std::pair<float, uint32_t>> candidates;
    while (!temp_queue.empty()) {
        candidates.emplace_back(temp_queue.top());
        temp_queue.pop();
    }

    for (auto& candidate : candidates) {
        uint32_t idx = candidate.second;
        float* vec = base + idx * vecdim;
        float dist = 0.0f;

        float32x4_t sum = vdupq_n_f32(0.0f);
        size_t d;
        for (d = 0; d + 4 <= vecdim; d += 4) {
            float32x4_t q = vld1q_f32(query + d);
            float32x4_t v = vld1q_f32(vec + d);
            float32x4_t diff = vsubq_f32(q, v);
            sum = vmlaq_f32(sum, diff, diff);
        }

        float32x2_t sum2 = vadd_f32(vget_low_f32(sum), vget_high_f32(sum));
        dist += vget_lane_f32(sum2, 0) + vget_lane_f32(sum2, 1);
        for (; d < vecdim; ++d) {
            float diff = query[d] - vec[d];
            dist += diff * diff;
        }
        dist = std::sqrt(dist);

        if (q.size() < k) {
            q.emplace(dist, idx);
        } else if (dist < q.top().first) {
            q.pop();
            q.emplace(dist, idx);
        }
    }

    return q;
}